#include <torch/torch.h>
#include <vector>
#include <cstdio>

at::Tensor gdw_conv_forward(
    at::Tensor x,
    at::Tensor w,
    at::Tensor b) {
    const auto batch_size = x.size(0);
    const auto in_channels = x.size(1);
    const auto height = x.size(2);
    const auto width = x.size(3);
    const auto out_channels = w.size(0);
    auto y = at::sum(at::sum(x * w.view({1, out_channels, height, width}), 3), 2) + b.view({1, -1});
    return y.view({batch_size, out_channels, 1, 1});
}

std::vector<at::Tensor> gdw_conv_backward(
    at::Tensor grad_y,
    at::Tensor x,
    at::Tensor w,
    at::Tensor b) {
    const auto batch_size = x.size(0);
    const auto in_channels = x.size(1);
    const auto height = x.size(2);
    const auto width = x.size(3);
    const auto out_channels = w.size(0);

    auto d_x = w.view({1, out_channels, height, width}) * grad_y;
    auto d_w = at::sum(x * grad_y, 0).unsqueeze(1);
    auto d_b = at::mean(grad_y, /*dim=*/0).squeeze();

    return {d_x, d_w, d_b};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &gdw_conv_forward, "global depthwise conv forward");
  m.def("backward", &gdw_conv_backward, "global depthwise conv backward");
}